import pandas as pd
from pyecharts.charts import Bar
from pyecharts import options as opts
import webbrowser

# 路径与年份
weather_path = "data/北京天气近十年_2015到2025.xlsx"
air_path = "data/北京空气质量近十年_2015到2025.xlsx"
sheet_years = [str(year) for year in range(2015, 2026)]

combined_data = []

for year in sheet_years:
    df_weather = pd.read_excel(weather_path, sheet_name=year)
    df_air = pd.read_excel(air_path, sheet_name=year)

    df_weather['日期'] = pd.to_datetime(df_weather['日期'], errors='coerce')
    df_air['日期'] = pd.to_datetime(df_air['日期'], errors='coerce')

    # 提取风向、风力等级
    df_weather['风向'] = df_weather['风力风向'].astype(str).str.extract(r'([东南西北]{1,2}风)')
    df_weather['风力等级'] = df_weather['风力风向'].astype(str).str.extract(r'(\d-\d)级')
    df_weather = df_weather[['日期', '风向', '风力等级']]

    df_air = df_air[['日期', 'PM2.5']]

    df_merge = pd.merge(df_weather, df_air, on='日期', how='inner')
    combined_data.append(df_merge)

# 合并并清洗
df_all = pd.concat(combined_data, ignore_index=True).dropna()
df_all = df_all[df_all['风向'].notnull() & df_all['风力等级'].notnull()]

# 分组计算平均 PM2.5
grouped = df_all.groupby(['风向', '风力等级'])['PM2.5'].mean().reset_index()

# 排序风向
wind_order = ['北风', '东北风', '东风', '东南风', '南风', '西南风', '西风', '西北风']
grouped['风向'] = pd.Categorical(grouped['风向'], categories=wind_order, ordered=True)
grouped = grouped.sort_values(['风向', '风力等级'])

# 准备数据
x_axis = sorted(grouped['风向'].dropna().unique())
wind_strengths = sorted(grouped['风力等级'].dropna().unique())

bar = Bar(init_opts=opts.InitOpts(width="1000px", height="600px"))
bar.add_xaxis(x_axis)

# 为每个风力等级添加一条系列
for strength in wind_strengths:
    data = []
    for wind in x_axis:
        match = grouped[(grouped['风向'] == wind) & (grouped['风力等级'] == strength)]
        val = round(match['PM2.5'].values[0], 2) if not match.empty else 0
        data.append(val)
    bar.add_yaxis(strength, data, stack="stack1")

# 设置图表样式
bar.set_global_opts(
    # title_opts=opts.TitleOpts(title="风向 + 风力等级对 PM2.5 的影响", pos_left="center"),
    xaxis_opts=opts.AxisOpts(name="风向"),
    yaxis_opts=opts.AxisOpts(name="平均PM2.5 (μg/m³)"),
    legend_opts=opts.LegendOpts(pos_top="5%"),
    tooltip_opts=opts.TooltipOpts(trigger="axis", axis_pointer_type="shadow")
)

bar.set_series_opts(label_opts=opts.LabelOpts(is_show=True, position="inside"))

# 输出并自动打开
output_file = "wind_pm25_bar_chart.html"
bar.render(output_file)
webbrowser.open(output_file)
